from typing import Dict, List, Set, Optional
from .CppCodeParser import CppCodeParser

class CallGraphAnalyzer:
    """Analyzes function call relationships in parsed C/C++ code."""
    
    def __init__(self, parser: CppCodeParser, debug_mode: bool = False):
        self.functions = parser.functions
        self.entry_points = parser.entry_points
        self.debug_mode = debug_mode
    
    def _debug_print(self, message: str):
        """Print a debug message if debug mode is enabled."""
        if self.debug_mode:
            print(f"[DEBUG] {message}")
    
    def build_call_tree(self, direction: str = "down", max_depth: int = None) -> Dict:
        """
        Build a call tree for all functions in the program.
        
        Args:
            direction: "down" for functions called by each function, "up" for functions that call each function
            max_depth: Maximum depth of the tree (None for unlimited)
        
        Returns:
            A dictionary representing the call tree for the entire program
        """
        tree = {}
        visited = set()
        
        def build_subtree(node, current_depth, tree_dict):
            if max_depth is not None and current_depth > max_depth:
                return
            
            if node.full_name in visited:
                tree_dict[f"{node.full_name} (recursive)"] = {}
                return
            
            visited.add(node.full_name)
            subtree = {}
            
            if direction == "down":
                # Make sure we have calls to process
                if node.calls:
                    for callee in node.calls:
                        build_subtree(callee, current_depth + 1, subtree)
                else:
                    # Debug: Print info when no calls are found
                    self._debug_print(f"No function calls found for {node.full_name}")
            elif direction == "up":
                if node.called_by:
                    for caller in node.called_by:
                        build_subtree(caller, current_depth + 1, subtree)
                else:
                    # Debug: Print info when no callers are found
                    self._debug_print(f"No callers found for {node.full_name}")
            
            visited.remove(node.full_name)
            tree_dict[node.full_name] = subtree
        
        # Build tree starting from each function
        for func_name, func in self.functions.items():
            visited = set()
            subtree = {}
            build_subtree(func, 1, subtree)
            tree[func_name] = subtree.get(func_name, {})
        
        return tree
    
    def generate_call_chains(self, target_name: str = None, max_depth: int = None) -> List[List[str]]:
        """
        Generate all possible function call chains in the program.
        
        Args:
            target_name: Optional target function to focus on
            max_depth: Maximum chain length (None for unlimited)
        
        Returns:
            List of call chains, each represented as a list of function names
        """
        chains = []
        visited = set()
        
        def dfs(node, path, depth):
            if max_depth is not None and depth > max_depth:
                return
            
            if node.full_name in visited:
                # We've found a recursive call, add the path and return
                chains.append(path + [f"{node.full_name} (recursive)"])
                return
            
            # If we've reached a leaf node (function that doesn't call others),
            # or if we've reached our target, add the current path to the chains
            if not node.calls or (target_name and node.full_name == target_name):
                chains.append(path)
                return
            
            visited.add(node.full_name)
            
            # For each function called by the current function, extend the path
            for callee in node.calls:
                dfs(callee, path + [callee.full_name], depth + 1)
            
            visited.remove(node.full_name)
        
        # Start from each entry point or from all functions if no entry points are defined
        start_points = [self.functions[name] for name in self.entry_points] if self.entry_points else self.functions.values()
        
        for start_node in start_points:
            visited = set()
            dfs(start_node, [start_node.full_name], 1)
        
        return chains
    
    def extract_computation_steps(self) -> Dict[str, str]:
        """
        Extract computation steps (function source code) for all functions in the call trees.
        
        Returns:
            Dictionary mapping function names to their source code
        """
        steps = {}
        
        for func_name, func in self.functions.items():
            if func.source_code:
                steps[func_name] = func.source_code
            else:
                steps[func_name] = f"// Source code not available for {func_name}"
        
        return steps 